package com.ftg.learn.controller;

import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.model.chat.ChatModel;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.util.Base64;

/**
 * @author 布衣
 */
@RestController
public class ImageToTextController {

    @Resource
    private ChatModel qwenVL;

    @PostMapping("/image-to-text")
    public String analyzeImage(@RequestParam("image") MultipartFile imageFile,
                               @RequestParam(value = "prompt", defaultValue = "请描述这张图片的内容") String prompt) {
        try {
            byte[] imageBytes = imageFile.getBytes();
            String base64Image = Base64.getEncoder().encodeToString(imageBytes);

            UserMessage userMessage = UserMessage.from(
                    TextContent.from(prompt),
                    ImageContent.from(base64Image, imageFile.getContentType())
            );

            String response = qwenVL.chat(userMessage).aiMessage().text();
            return response;

        } catch (Exception e) {
            return "分析图片时发生错误: " + e.getMessage();
        }
    }

    @PostMapping("/image-description")
    public String describeImage(@RequestParam("image") MultipartFile imageFile) {
        try {
            byte[] imageBytes = imageFile.getBytes();
            String base64Image = Base64.getEncoder().encodeToString(imageBytes);

            UserMessage userMessage = UserMessage.from(
                    TextContent.from("请详细描述这张图片的内容，包括图片中的主要对象、场景和活动。"),
                    ImageContent.from(base64Image, imageFile.getContentType())
            );

            String response = qwenVL.chat(userMessage).aiMessage().text();
            return response;

        } catch (Exception e) {
            return "描述图片时发生错误: " + e.getMessage();
        }
    }

    @PostMapping("/ocr")
    public String ocrImage(@RequestParam("image") MultipartFile imageFile) {
        try {
            byte[] imageBytes = imageFile.getBytes();
            String base64Image = Base64.getEncoder().encodeToString(imageBytes);

            UserMessage userMessage = UserMessage.from(
                    TextContent.from("请识别并提取这张图片中的所有文字内容。按原文格式输出。"),
                    ImageContent.from(base64Image, imageFile.getContentType())
            );

            String response = qwenVL.chat(userMessage).aiMessage().text();
            return response;

        } catch (Exception e) {
            return "OCR识别时发生错误: " + e.getMessage();
        }
    }
}

